package com.aryankh.core;

import com.aryankh.utils.XssUtil;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.util.HashMap;
import java.util.Map;

/**
 * @version 1.0
 * @Author AryanKH
 * @Date 2022/6/6 14:37
 * @Default
 */
public class XssHttpServletRequestWapper extends HttpServletRequestWrapper {
    private final HttpServletRequest orgRequest;

    public XssHttpServletRequestWapper(HttpServletRequest request) {
        super(request);
        orgRequest=request;
    }

    //重写这几个方法 并且处理xss相关内容
    @Override
    public String getParameter(String name) {
        return XssUtil.xssEncode(super.getParameter(XssUtil.xssEncode(name)));
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        Map<String, String[]> parameterMap = super.getParameterMap();
        Map<String, String[]> newParameterMap = new HashMap<>(8);
        for(Map.Entry<String, String[]> entry: parameterMap.entrySet()){
            //处理key的标签
            String key = entry.getKey();
            String newKey = XssUtil.xssEncode(key);
            String[] params = entry.getValue();
            String[] newParams = new String[params.length];
            for (int i = 0; i < params.length; i++) {
                newParams[i] = XssUtil.xssEncode(params[i]);
            }
            newParameterMap.put(newKey, newParams);
        }
        return newParameterMap;
    }

    @Override
    public String[] getParameterValues(String name) {
        String[] parameterValues = super.getParameterValues(name);
        String[] newParamsValues = new String[parameterValues.length];
        for (int i = 0; i < parameterValues.length; i++) {
            newParamsValues[i] = XssUtil.xssEncode(parameterValues[i]);
        }
        return newParamsValues;
    }

    // 如果将来想继续使用原始的request可以通过这个方法获取
    public HttpServletRequest getOrgRequest() {
        return orgRequest;
    }
}
